package com.dong.sharding.config;

import com.google.common.collect.Range;
import org.apache.shardingsphere.sharding.api.sharding.standard.PreciseShardingValue;
import org.apache.shardingsphere.sharding.api.sharding.standard.RangeShardingValue;
import org.apache.shardingsphere.sharding.api.sharding.standard.StandardShardingAlgorithm;
import org.springframework.stereotype.Component;

import java.text.SimpleDateFormat;
import java.time.LocalDateTime;
import java.time.ZoneId;
import java.time.format.DateTimeFormatter;
import java.util.Collection;
import java.util.Date;
import java.util.LinkedHashSet;

/**
 * 自定义分片算法-时间分片算法
 *
 * @author xiedongxiao
 */

@Component
public class TimeBasedShardingAlgorithm implements StandardShardingAlgorithm<Date> {
    @Override
    public String doSharding(Collection<String> availableTargetNames, PreciseShardingValue<Date> shardingValue) {
        Date value = shardingValue.getValue();
        SimpleDateFormat sdf = new SimpleDateFormat("yyyyMM");
        String shardKey = sdf.format(value);

        // 假设你的数据源或表名按照 yyyyMM 的方式命名，例如：log_202505
        String targetTableName = "log_" + shardKey;

        if (availableTargetNames.contains(targetTableName)) {
            return targetTableName;
        } else {
            // 如果找不到匹配的分片，可以选择默认分片或者抛出异常
            throw new IllegalArgumentException("No matching data source found for [" + shardKey + "]");
        }
    }

    @Override
    public Collection<String> doSharding(Collection<String> availableTargetNames, RangeShardingValue<Date> rangeShardingValue) {
        Collection<String> result = new LinkedHashSet<>();

        Range<Date> range = rangeShardingValue.getValueRange();
        LocalDateTime startDateTime = range.lowerEndpoint().toInstant()
                .atZone(ZoneId.systemDefault())
                .toLocalDateTime();
        LocalDateTime endDateTime = range.upperEndpoint().toInstant()
                .atZone(ZoneId.systemDefault())
                .toLocalDateTime();

        for (LocalDateTime date = startDateTime; !date.isAfter(endDateTime); date = date.plusMonths(1)) {
            String formattedDate = date.format(DateTimeFormatter.ofPattern("yyyyMM"));
            String targetTableName = rangeShardingValue.getLogicTableName() + "_" + formattedDate;
            if (availableTargetNames.contains(targetTableName)) {
                result.add(targetTableName);
            } else {
                throw new IllegalArgumentException("Table " + targetTableName + " does not exist.");
            }
        }
        return result;
    }
}
